# -*- coding: utf-8 -*-
import logging
from ctypes import *
import traceback
import ipaddr
import os
import re

#
# reference: //Core/NCSG/Dev/TDA-3.8SP3/TDA/src/trend/cav/src/dtas_feedback.c ,RuleManager.java
#


class _NS_RULE_SET(Structure): pass
NS_RULE_SET = _NS_RULE_SET
NS_RULE_SET_HANDLE = POINTER(NS_RULE_SET)


class _NS_RULE_GROUP(Structure): pass
NS_RULE_GROUP = _NS_RULE_GROUP
NS_RULE_GROUP_HANDLE = POINTER(NS_RULE_GROUP)


class NS_RULE_MATCH_RESULT(Structure):
    _fields_ = [("priority", c_uint32), ("action", c_uint32), ("ruleId", c_uint32)]

class NS_IP_ADDRESS_T(Union):
    _fields_ = [('v4_byte', c_uint8 * 4),
                ('v4_addr16', c_int16 * 2),
                ('v4', c_int32),
                ('v6_byte', c_uint8 * 16),
                ('v6_addr16', c_int16 * 8),
                ('v6', c_int32 * 4)]

LIBRULE_FILE='/opt/TrendMicro/DTAS/ManagementServer/lib/librule.so'




NS_RULE_ACTION_BLOCK = 0x1001

class Type(object):
    IP = 0
    URL = 1
    SHA1 = 2
    DOMAIN = 3

class ruleException(Exception):pass

class LibruleWrapper(object):

    def __init__(self, logger = None):
        if logger:
            self.log = logger
        else:
            self.log = logging

        if os.path.exists(LIBRULE_FILE):
            self._init_librule(LIBRULE_FILE)
        else:
            self.log.error('Fail to find {}'.format(LIBRULE_FILE))

        self.group_handler_url = NS_RULE_GROUP_HANDLE()
        self.group_handler_domain = NS_RULE_GROUP_HANDLE()
        self.group_handler_ip = NS_RULE_GROUP_HANDLE()

        self.group_hander_sha1 = []

        # succeed to matched
        self._Matched = 1

        # valid items
        self._Valid = 0

        self.sha1_pattern = re.compile(r'^[0-9a-fA-F]{40}$')
        self.has_ipv6_pattern = re.compile(r'.*\[(.*)\].*')
        self.url_valid_pattern = re.compile(r'^(http:\/\/|https:\/\/)*(\*|([A-Z0-9][A-Z0-9_-]*)+)(\.[A-Z0-9][A-Z0-9_-]*)+(\:\d{1,5})?((\/(\*|$|([a-zA-Z0-9\.\,\;\?\'\\\+\{\}&\$#\=~_\-]|(%[A-F0-9]{2}))+|))*)',re.I)
        self.domain_pattern = re.compile(r'^([a-z\d](-*[a-z\d])*)(\.([a-z\d](-*[a-z\d])*))*$',re.I)

        self.ip_range=[]

    def _init_librule(self,librule_path):

        self._librule = cdll.LoadLibrary(librule_path)
        self._librule.ns_validateRule_IP.argtypes = [c_char_p, c_uint32]
        self._librule.ns_validateRule_IP.restype = c_int
        self._librule.ns_validateRule_FQDN.argtypes = [c_char_p, c_uint32]
        self._librule.ns_validateRule_FQDN.restype = c_int
        self._librule.ns_validateRule_URL.argtypes = [c_char_p, c_uint32]
        self._librule.ns_validateRule_URL.restype = c_int
        self._librule.ns_newRuleSet.argtypes = [c_void_p, c_void_p]
        self._librule.ns_newRuleSet.restype = c_int
        self._librule.ns_delRuleSet.argtypes = [c_void_p]
        self._librule.ns_delRuleSet.restype = c_int
        self._librule.ns_addRuleToRuleSet.argtypes = [c_void_p, c_char_p, c_uint32, c_uint32, c_uint32]
        self._librule.ns_addRuleToRuleSet.restype = c_int
        self._librule.ns_buildRuleGroup.argtypes = [c_void_p, c_void_p]
        self._librule.ns_buildRuleGroup.restype = c_int
        self._librule.ns_delRuleGroup.argtypes = [c_void_p]
        self._librule.ns_delRuleGroup.restype = c_int
        self._librule.ns_matchIPRuleGroup.argtypes = [c_void_p, c_uint32, c_void_p]
        self._librule.ns_matchIPRuleGroup.restype = c_int
        self._librule.ns_matchIPRuleGroupEx.argtypes = [c_void_p, c_uint32, c_void_p, c_void_p]
        self._librule.ns_matchIPRuleGroupEx.restype = c_int
        self._librule.ns_matchUrlWithRuleGroup.argtypes = [c_void_p, c_char_p, c_uint32, c_void_p]
        self._librule.ns_matchUrlWithRuleGroup.restype = c_int
        self._librule.ns_matchDomainNameWithRuleGroup.argtypes = [c_void_p, c_char_p, c_uint32, c_void_p]
        self._librule.ns_matchDomainNameWithRuleGroup.restype = c_int

    def __del__(self):
        self.reset()

    def reset(self):
        try:
            if self.group_handler_url:
                self._librule.ns_delRuleGroup(byref(self.group_handler_url))
            if self.group_handler_domain:
                self._librule.ns_delRuleGroup(byref(self.group_handler_domain))
            if self.group_handler_ip:
                self._librule.ns_delRuleGroup(byref(self.group_handler_ip))

            del self.group_hander_sha1[:]
        except Exception:
            self.log.debug("Release allcated librule resource failed! {}".format(traceback.format_exc()))

    def _new_rule_set(self,rule_set_type):

        handler = NS_RULE_SET_HANDLE()
        if self._librule.ns_newRuleSet(rule_set_type,byref(handler)) != 0:
            self.log.error("ERROR: failed to build newRuleSet")
            return None
        return handler

    def _build_rule_group(self,rule_set_type):

        group_handler = NS_RULE_GROUP_HANDLE()
        if self._librule.ns_buildRuleGroup(rule_set_type, byref(group_handler)) !=0:
            self.log.error('Fail to build rule group')
            return None
        return group_handler

    def _pre_process(self,type,data):

        if type == Type.URL:
            # strip query string (?) and pound sign (#)
            idx1 = data.find('?')
            if (idx1 > 0):
                idx2 = data[:idx1].find('#')
                if (idx2 > 0):
                    data = data[:idx2]
                else:
                    data = data[:idx1]
            else:
                idx2 = data.find('#')
                if (idx2 > 0):
                    data = data[:idx2]
                else:
                    data = data
        elif type == Type.IP and data.find('-') != -1:
            self.ip_range.append(data)
        return data

    def import_rules(self,rules):
        '''
        format:{ ruleType1:[data1,data2,...],ruleType2:[data3,data4,...]}
        :param rules:
        :return:
        '''

        NS_RULE_SET_TYPE_IP = 1
        NS_RULE_SET_TYPE_FQDN = 2
        NS_RULE_SET_TYPE_URL = 3

        if not isinstance(rules,dict):
            self.log.error('Input is not a dict')
            return False

        rule_set_domain = self._new_rule_set(NS_RULE_SET_TYPE_FQDN)
        rule_set_ip = self._new_rule_set(NS_RULE_SET_TYPE_IP)
        rule_set_url = self._new_rule_set(NS_RULE_SET_TYPE_URL)

        rule_id = 1

        type_2_handle = {
            Type.DOMAIN:rule_set_domain,
            Type.IP:rule_set_ip,
            Type.URL:rule_set_url
        }
        for rule_type,rule_data in rules.iteritems():

            if not isinstance(rule_data,list):
                raise ruleException('Rule data must be a list object [{}]'.format(rule_data))

            for data in rule_data:
                if rule_type == Type.SHA1:
                    self.group_hander_sha1.append(str(data).upper())
                    continue

                rule_handle = type_2_handle.get(rule_type,None)

                if rule_handle is None:
                    raise ruleException('No such type [{}]'.format(rule_type))

                data = self._pre_process(rule_type,data)

                if self._librule.ns_addRuleToRuleSet(rule_handle,data,len(data),NS_RULE_ACTION_BLOCK,rule_id) != 0:
                    self.log.error('Librule ns_addRuleToRuleSet error {}'.format(data))
                    continue

                rule_id += 1

        self.group_handler_domain = self._build_rule_group(rule_set_domain)
        self.group_handler_ip = self._build_rule_group(rule_set_ip)
        self.group_handler_url = self._build_rule_group(rule_set_url)

        # clean rule set handler
        self._librule.ns_delRuleSet(byref(rule_set_ip))
        self._librule.ns_delRuleSet(byref(rule_set_ip))
        self._librule.ns_delRuleSet(byref(rule_set_url))

    def match_domain(self,domain):

        rule_match_result = NS_RULE_MATCH_RESULT()
        result = self._librule.ns_matchDomainNameWithRuleGroup(self.group_handler_domain,domain,len(domain),byref(rule_match_result))
        logging.debug('matchDomain:[{}] result {}'.format(domain,result))
        return result == self._Matched

    def _match_ipv6(self,ip_obj):

        uaddr = NS_IP_ADDRESS_T()

        ipv6_bytes = ip_obj.packed
        for i in range(len(ipv6_bytes)):
            uaddr.v6_byte[i] = ord(ipv6_bytes[i])

        rule_match_result = NS_RULE_MATCH_RESULT()
        result = self._librule.ns_matchIPRuleGroupEx(self.group_handler_ip,6,byref(uaddr),byref(rule_match_result))

        logging.debug('matchIPv6:[{}] result {}'.format(str(ip_obj),result))
        return result == self._Matched

    def _match_ipv4(self,ip_obj):

        rule_match_result = NS_RULE_MATCH_RESULT()
        result = self._librule.ns_matchIPRuleGroup(self.group_handler_ip,int(ip_obj),byref(rule_match_result))
        logging.debug('matchIPv4:[{}] result {}'.format(str(ip_obj),result))
        return result == self._Matched

    def _match_ip(self,ip):
        '''
        (single IP): <Digit>.<Digit>.<Digit>.<Digit>
        (subnet): <Digit>.<Digit>.<Digit>.<Digit>/<Mask Bit>
        Support match IP or IP subnet
        :param ip:
        :return:
        '''
        subnet = ip.split('/')
        if len(subnet) == 2:
            logging.debug('IP subnet find')
            ipA,mask = subnet
        else:
            logging.debug('IP single find')
            ipA,mask = ip,None
        if ipaddr.IPv6Address == type(ipaddr.IPAddress(ipA)):
            ip_v = 6
            if mask:
                ip_obj = ipaddr.IPv6Network(ip)
            else:
                ip_obj = ipaddr.IPv6Address(ip)
        elif ipaddr.IPv4Address == type(ipaddr.IPAddress(ipA)):
            ip_v = 4
            if mask:
                ip_obj = ipaddr.IPv4Network(ip)
            else:
                ip_obj = ipaddr.IPv4Address(ip)

        if ip_v == 6:
            return self._match_ipv6(ip_obj)
        else:
            return self._match_ipv4(ip_obj)

    def _match_ip_range(self,range):
        '''
        (IP range): <Digit>.<Digit>.<Digit>.<Digit>-<Digit>.<Digit>.<Digit>.<Digit>
        Support match IP range
        :param ip:
        :return:
        '''

        r = range.split('-')
        for i in ipaddr.summarize_address_range(ipaddr.IPAddress(r[0]),ipaddr.IPAddress(r[1])):
            logging.debug('Test subnet {}'.format(str(i)))
            if self._match_ip(str(i)):
                return True
        return False

    def _match_ip_range_str(self,range):
        '''
        Only check if IP range is match  with string comparison
        :param range:
        :return:
        '''
        return range in self.ip_range

    def match_ip(self,ip):
        '''
        Support IP,IP subnet,IP range
        :param ip:
        :return:
        '''
        if ip.find('-') != -1:
            logging.debug('IP range find')
            # if ip range, do not check duplicate
            return self._match_ip_range_str(ip)
            #return self._match_ip_range(ip)
        else:
            return self._match_ip(ip)

    def match_url(self,url):
        rule_match_result = NS_RULE_MATCH_RESULT()
        result = self._librule.ns_matchUrlWithRuleGroup(self.group_handler_url, url,len(url), byref(rule_match_result))
        logging.debug('matchurl:[{}] result {}'.format(url,result))
        return result == self._Matched

    def match_sha1(self,sha1):
        return str(sha1).upper() in self.group_hander_sha1

    def match_all(self,data,type):
        if type == Type.DOMAIN:
            return self.match_domain(data)
        elif type == Type.URL:
            return self.match_url(data)
        elif type == Type.IP:
            return self.match_ip(data)
        elif type == Type.SHA1:
            return self.match_sha1(data)
        else:
            raise ruleException('No such type [{}]'.format(type))

    def _valide_url_byself(self,url):

        if len(url) > 2138 or len(url)<1:
            return False
        # check if has ipv6 address in url
        if self.has_ipv6_pattern.match(url):
            url=re.sub(r'\[.*\]','1.1.1.1',url)
        if self.url_valid_pattern.match(url):
            return True
        else:
            logging.debug('url:[{}] is invalid'.format(url))
            return False

    def is_valid_url(self,url):
        url = url.strip()
        if not self._valide_url_byself(url):
            return False

        result = self._librule.ns_validateRule_URL(url, len(url))
        logging.debug('validateURL:[{}],result {}'.format(url,result))
        return result == self._Valid

    def is_valid_ip(self,ip):
        result = self._librule.ns_validateRule_IP(ip, len(ip))
        logging.debug('validateIP:[{}] result {}'.format(ip,result))
        return result == self._Valid

    def _is_valid_domain_byself(self,domain):

         if len(domain) >253:
             return False
         for i in domain.split('.'):
             if len(i)> 63:
                 return False
         if self.domain_pattern.match(domain):
             return True
         else:
             logging.debug('domain: [{}] is invalid'.format(domain))
             return False

    def is_valid_domain(self,domain):
        domain = domain.strip()
        if not self._is_valid_domain_byself(domain):
            return False
        result = self._librule.ns_validateRule_FQDN(domain, len(domain))
        logging.debug('validateDomain:[{}] result {}'.format(domain,result))
        return result == self._Valid

    def is_valid_sha1(self,sha1):
        result = self.sha1_pattern.match(sha1) is not None
        logging.debug('validateSha1:[{}] result {}'.format(sha1,result))
        return result

    def is_valid_all(self,data,type):
        if type == Type.DOMAIN:
            return self.is_valid_domain(data)
        elif type == Type.URL:
            return self.is_valid_url(data)
        elif type == Type.IP:
            return self.is_valid_ip(data)
        elif type == Type.SHA1:
            return self.is_valid_sha1(data)
        else:
            raise ruleException('No such type [{}]'.format(type))

if __name__ == '__main__':
    import sys
    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)

    librule = LibruleWrapper()

    d1='baidu.com'
    d2='www.sina.com'

    i1='1.1.1.1'
    i2='fe80::480e:90ff:fe9f:60f3'
    i3='fe80::480e:90ff:fe9f:60f2'


    u1 = 'http://www.baidu.com/'
    u2 = 'http://www.163.com/#123?a=1&b=2'

    s1 = '4437E47D078A6582BF7EF4784925B4E13A129EAB'
    s2 = '0A295CB460193F2CE1FBF58FF83BE130142CDF5B'

    rules={}
    rules[Type.URL] = [u1,u2]
    rules[Type.DOMAIN] = [d1,d2]
    rules[Type.IP] = [i1,i2,i3]
    rules[Type.SHA1]=[s1,s2]

    librule.import_rules(rules)

    print 'Test Match OK...'
    print librule.match_domain(d1)
    print librule.match_domain(d2)
    print librule.match_ip(i1)
    print librule.match_ip(i2)
    print librule.match_ip(i3)
    print librule.match_url(u1)
    print librule.match_url(u2)
    print librule.match_sha1(s1)
    print librule.match_sha1(s2)

    print 'Test Match Wrong ...'
    print librule.match_domain('stz.com')
    print librule.match_ip('6.6.6.6')
    print librule.match_ip('fe80::20c:29ff:fece:929b')
    print librule.match_url('http://www.stz.com/#123?a=1&b=2')
    print librule.match_sha1('44')

    print 'Test After Reset Match ...'
    librule.reset()
    print librule.match_domain(d1)
    print librule.match_domain(d2)
    print librule.match_ip(i1)
    print librule.match_ip(i2)
    print librule.match_ip(i3)
    print librule.match_url(u1)
    print librule.match_url(u2)
    print librule.match_sha1(s1)
    print librule.match_sha1(s2)

    print 'Test Append source OK Match ...'

    rules[Type.URL]=[u1]
    rules[Type.DOMAIN] = [d1]
    rules[Type.IP] = [i1]
    rules[Type.SHA1]=[s1]

    librule.import_rules(rules)

    rules[Type.URL]=[u2]
    rules[Type.DOMAIN] = [d2]
    rules[Type.IP] = [i2,i3]
    rules[Type.SHA1]=[s2]

    librule.import_rules(rules)
    print '[fist] ...'
    print librule.match_url(u1)
    print librule.match_domain(d1)
    print librule.match_ip(i1)
    print librule.match_sha1(s1)

    print '[Second]...'
    print librule.match_url(u2)
    print librule.match_domain(d2)
    print librule.match_ip(i2)
    print librule.match_ip(i3)
    print librule.match_sha1(s2)
    librule.reset()
    print '[Match test: Only SHA1 support the append,and others do not support the appand]'


    print 'Test valid OK...'
    print librule.is_valid_domain('baidu.com')
    print librule.is_valid_url('http://baidu.com')
    print librule.is_valid_ip('1.1.1.1')
    print librule.is_valid_ip('fe80::480e:90ff:fe9f:60f2')
    print librule.is_valid_sha1('94461E02B02B9D0B27664A7A8EF003AEE3CAFFaA')

    print 'Test valid Wrong...'
    print librule.is_valid_domain('#')
    print librule.is_valid_url('http:baiducom')
    print librule.is_valid_ip('1.1.1.1999')
    print librule.is_valid_ip('fe80::480e:90ff:')
    print librule.is_valid_sha1('94461E02B02B9D0B27664A7A8EF003AEE3CAFFa')

    print 'Test valid IP range/subnet OK...'
    print librule.is_valid_ip('5.5.5.5/5')
    print librule.is_valid_ip('2001:db8:abcd:1000::/52')
    print librule.is_valid_ip('1.1.1.1-1.1.1.25')
    print librule.is_valid_ip('fdda:6d07:b33c:49e7::-fdda:6d07:b33c:49e7:ffff:ffff:ffff:ffff')

    print 'Test valid IP range/subnet Wrong...'
    print librule.is_valid_ip('5.5.5.5/56')
    print librule.is_valid_ip('1.1.1.1-1.1.1.2599')


    print 'Test Match IP range/subnet OK...'
    rules={}
    rules[Type.IP] = ['5.5.5.5/5','2001:db8:abcd:1000::/52','1.1.1.1-1.1.1.25','1.1.1.1-1.1.1.2599',
                      'fdda:6d07:b33c:49e7::-fdda:6d07:b33c:49e7:ffff:ffff:ffff:ffff']
    librule.import_rules(rules)
    print librule.match_ip('1.1.1.1-1.1.1.25')
    print librule.match_ip('5.5.5.25')
    print librule.match_ip('5.5.5.5/5')
    print librule.match_ip('2001:db8:abcd:1000::/52')


    print 'Test Match IP range/subnet Wrong...'
    print librule.match_ip('20.20.20.20/20')
    print librule.match_ip('10.1.1.1-10.1.1.25')
    print librule.match_ip('1001:db8:abcd:1000::/22')
    print librule.match_ip('edda:6d07:b33c:49e7::-fdda:6d07:b33c:49e7::ffff:ffff:ffff')

    librule.reset()

    print 'Self check ...'
    print librule.is_valid_url('http:baidu.com')
    print librule.is_valid_url('http://baidu.com')
    print librule.is_valid_url('http//test.com')
    print librule.is_valid_domain('baidu')
    print librule.is_valid_domain('baidu.com')
